import cv2
import pandas as pd
import numpy as np
import yaml

# 配置参数
img_path = '../maps/map2.png'
traj_path = '../maps/trajectory.csv'
output_path = '../maps/map2_with_traj.png'

# 手动修正参数（世界系下，轨迹点需要先逆向应用此修正）
manual_offset = [0,0,0]  # [x, y, yaw]，单位: 米, 米, 弧度
# 例如: manual_offset = [0.1, -0.2, 0.05]

# 加载图像
img = cv2.imread(img_path)
if img is None:
    raise FileNotFoundError(f"图像文件未找到: {img_path}")

# 加载轨迹
df = pd.read_csv(traj_path)

# 读取地图参数
yaml_path = '../maps/map2.yaml'
with open(yaml_path, 'r', encoding='utf-8') as f:
    map_cfg = yaml.safe_load(f)
resolution = map_cfg['resolution']  # 米/像素
origin = map_cfg['origin']          # [origin_x, origin_y, origin_yaw]
origin_x, origin_y, origin_yaw = origin[0], origin[1], origin[2]

# 轨迹点转换为像素坐标（先手动修正，再考虑地图yaw旋转）
pts = []
mx, my, myaw = manual_offset
cos_myaw = np.cos(-myaw)
sin_myaw = np.sin(-myaw)
cos_yaw = np.cos(-origin_yaw)
sin_yaw = np.sin(-origin_yaw)
for _, row in df.iterrows():
    # 1. 先对轨迹点做逆向旋转和平移（手动修正）
    tx = row['x'] - mx
    ty = row['y'] - my
    rx = tx * cos_myaw - ty * sin_myaw
    ry = tx * sin_myaw + ty * cos_myaw
    # 2. 再用地图参数做旋转和平移
    dx = rx - origin_x
    dy = ry - origin_y
    rrx = dx * cos_yaw - dy * sin_yaw
    rry = dx * sin_yaw + dy * cos_yaw
    # 3. 转为像素坐标
    px = int(rrx / resolution)
    py = img.shape[0] - int(rry / resolution)
    pts.append([px, py])
pts = np.array(pts, dtype=np.int32)

# 绘制轨迹
for i in range(1, len(pts)):
    cv2.line(img, tuple(pts[i-1]), tuple(pts[i]), (0,0,0), 3)  

# 计算世界系原点在像素坐标中的位置
# 世界原点 [0,0] 经过地图参数和手动修正后，映射到像素坐标
def world_to_pixel(wx, wy):
    # 手动修正
    tx = wx - mx
    ty = wy - my
    rx = tx * np.cos(-myaw) - ty * np.sin(-myaw)
    ry = tx * np.sin(-myaw) + ty * np.cos(-myaw)
    # 地图参数
    dx = rx - origin_x
    dy = ry - origin_y
    rrx = dx * np.cos(-origin_yaw) - dy * np.sin(-origin_yaw)
    rry = dx * np.sin(-origin_yaw) + dy * np.cos(-origin_yaw)
    px = int(rrx / resolution)
    py = img.shape[0] - int(rry / resolution)
    return px, py

# 绘制原点坐标架
origin_px, origin_py = world_to_pixel(0, 0)
axis_len = int(0.5 / resolution)  # 0.5米长度的箭头
# x轴终点（世界系x正方向0.5米）
x_end = world_to_pixel(2, 0)
# y轴终点（世界系y正方向0.5米）
y_end = world_to_pixel(0, 2)
# 绘制x轴
cv2.arrowedLine(img, (origin_px, origin_py), x_end, (255,0,0), 2, tipLength=0.15)
# 绘制y轴
cv2.arrowedLine(img, (origin_px, origin_py), y_end, (0,255,0), 2, tipLength=0.15)

# 保存结果
cv2.imwrite(output_path, img)
print(f"轨迹已绘制并保存到: {output_path}")
